import numpy as np
import trimesh
import os

def depth_to_point_cloud(depth_image, intrinsics):
    """
    将深度图像转换为点云。
    
    参数:
        depth_image (ndarray): 形状为 [H, W] 的深度图像，单位为米或毫米。
        intrinsics (dict): 相机内参字典，包含 'fx', 'fy', 'cx', 'cy'（焦距和光心）。
    
    返回:
        point_cloud (trimesh.points.PointCloud): Trimesh 点云对象。
    """
    H, W = depth_image.shape
    fx, fy, cx, cy = intrinsics['fx'], intrinsics['fy'], intrinsics['cx'], intrinsics['cy']
    
    # 创建网格坐标（像素坐标）
    u, v = np.meshgrid(np.arange(W), np.arange(H))  # (v, u) = (row, col)
    
    # 根据深度图计算三维空间坐标 (x, y, z)
    Z = depth_image[v, u]  # 深度值
    X = (u - cx) * Z / fx  # 转换为三维坐标 x
    Y = (v - cy) * Z / fy  # 转换为三维坐标 y
    
    # 创建点云数据
    points = np.vstack((X.ravel(), Y.ravel(), Z.ravel())).T  # [N, 3]
    
    # 使用 trimesh 创建点云对象
    point_cloud = trimesh.points.PointCloud(points)
    
    return point_cloud

dir = 'sample_out/depth_npy'

if __name__ == '__main__':
    # 相机内参
    intrinsics = {
                'fx': 500,  # 焦距x
                'fy': 500,  # 焦距y
                'cx': 320,  # 光心x
                'cy': 240   # 光心y
            }
    
    files = os.listdir(dir)
    for file in files:
        if file.endswith('.npy'):
            print(f"Processing {file}...")
            depth_image = np.load(f"{dir}/{file}")
            point_cloud = depth_to_point_cloud(depth_image, intrinsics)
            point_cloud.export(f"{dir}/{file[:-4]}.ply")
            print(f"Saved {file[:-4]}.ply")
    print("Done.")